"""
自定义一个本地的Embeddings
"""
from chromadb import Embeddings
from sentence_transformers import SentenceTransformer


class BGE_Embed(Embeddings):

    def __init__(self):
        self.model = SentenceTransformer('./data/BAAI/bge-small-zh')

    # 对文档进行词嵌入
    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        return [self.model.encode(text, normalize_embeddings=True).tolist() for text in texts]

    # 对查询语句进行嵌入
    def embed_query(self, text: str) -> list[float]:
        return self.model.encode(text, normalize_embeddings=True).tolist()